"""
@Time: 2021/1/24 下午 9:53
@Author: jinzhuan
@File: ee_toolkit.py
@Desc: 
"""
import torch
torch.cuda.set_device(4)

from cognlp import *
from cognlp.toolkit.ee.ee_toolkit import EeToolkit

tokenize_toolkit = TokenizeToolkit()
words = tokenize_toolkit.run(
    "And these bozos let four armed Cubans land on our shores when they're trying to make a high terrorist alert.")

ner_toolkit = NerAce2005Toolkit(
    bert_model='bert-base-cased',
    model_path='../../../cognlp/data/ner/ace2005/toolkit/model.pt',
    vocabulary_path='../../../cognlp/data/ner/ace2005/toolkit/vocabulary.txt',
    device=torch.device("cuda"),
    device_ids=[4],
    max_seq_length=256)
ner_result = ner_toolkit.run(words)

ee_toolkit = EeToolkit(model_path='../../../cognlp/data/ee/ace2005/toolkit/model.pt',
                       trigger_vocabulary_path='../../../cognlp/data/ee/ace2005/toolkit/vocabulary.txt',
                       argument_vocabulary_path='../../../cognlp/data/ee/ace2005/toolkit/argument_vocabulary.txt',
                       device_ids=[4],
                       max_seq_length=256)
ee_result = ee_toolkit.run(words, ner_result)
print(ee_result)
